import matplotlib.pyplot as plt
import os
from collections import defaultdict

def plot_metrics(history, is_show=True, save_path=None):
    metrics = ['loss', 'accuracy', 'precision', 'recall']
    if is_show or save_path is not None:
        plt.figure(figsize=(20, 4))  # 宽度适配4张图

        for i, metric in enumerate(metrics):
            plt.subplot(1, 4, i + 1)
            plt.plot(history.history[metric], label=f'Train {metric}')
            plt.plot(history.history[f'val_{metric}'], label=f'Val {metric}')
            plt.title(f'{metric.capitalize()} over Epochs')
            plt.xlabel('Epochs')
            plt.ylabel(metric.capitalize())
            plt.legend()
            plt.grid(True)

        plt.tight_layout()
        if save_path is not None and is_show:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            plt.savefig(save_path)
            plt.show()
        elif save_path:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            plt.savefig(save_path)
            plt.close()
        elif is_show:
            plt.show()
def plot_metrics_multiple(history, is_show=True, save_path=None):
    # 自动提取所有指标名称
    all_keys = list(history.history.keys())
    grouped_metrics = defaultdict(dict)

    # 归类为 train/val
    for key in all_keys:
        if key.startswith('val_'):
            metric = key[4:]
            grouped_metrics[metric]['val'] = history.history[key]
        else:
            grouped_metrics[key]['train'] = history.history[key]

    # 每个指标单独绘图
    num_metrics = len(grouped_metrics)
    plt.figure(figsize=(5 * num_metrics, 4))

    for i, (metric, values) in enumerate(grouped_metrics.items()):
        plt.subplot(1, num_metrics, i + 1)
        if 'train' in values:
            plt.plot(values['train'], label='Train')
        if 'val' in values:
            plt.plot(values['val'], label='Val')
        plt.title(f'{metric.capitalize()}')
        plt.xlabel('Epochs')
        plt.ylabel(metric.capitalize())
        plt.legend()
        plt.grid(True)

    plt.tight_layout()

    if save_path is not None:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        plt.savefig(save_path)
    if is_show:
        plt.show()
    else:
        plt.close()
